import numpy as np

class ZF_SICDetector(object):

    def __init__(self, params):

        self.Nr = params['Nr']
        self.Nt = params['Nt']
        self.batch_size = params['BATCH_SIZE']
        self.C = np.expand_dims(params['CONSTELLATION_COMPLEX'],0) # [1,M]

    def __call__(self,x,H,y,sigma2):
        """
        迫零干扰消除检测算法
        Input:
        x:发送符号 array(shape=(batch_size, 2*Nt), dtype=float32)
        y:接收信号 array(shape=(batch_size, 2*Nr), dtype=float32)
        H:信道矩阵 array(shape=(batch_size, 2*Nr, 2*Nt), dtype=float32)
        sigma2:噪声方差 array(shape=(batch_size, 1), dtype=float32)
        Output:
        nodes:指标字典
        """
        # 实复变换
        H_c = H[:,0:self.Nr,0:self.Nt] + 1j*H[:,self.Nr:self.Nr*2,0:self.Nt]
        y_c = y[:,0:self.Nr] + 1j*y[:,self.Nr:self.Nr*2]
        x_c = x[:,0:self.Nt] + 1j*x[:,self.Nt:self.Nt*2]

        xhat_ = np.zeros([self.batch_size,self.Nt]).astype('complex128')
        index = list(range(self.batch_size))
        ktrans = np.reshape(range(self.Nt),[1,self.Nt])
        ktrans = np.tile(ktrans,[self.batch_size,1])

        for i in range(self.Nt):

            # 检测
            HtH = np.matmul(np.conjugate(np.transpose(H_c,[0, 2, 1])), H_c)
            HtHinv = np.linalg.inv(HtH)
            Q = np.matmul(HtHinv, np.conjugate(np.transpose(H_c,[0, 2, 1]))) # [b,Nt,Nr]
            xhat = self.batch_matvec_mul(Q, y_c)
            xhat = self.slicer_c(xhat)

            # 选择
            p = np.sum(np.square(np.abs(Q)),2) # [b,Nt]
            k = np.argmin(p,axis=1) # [b,]
            k1 = ktrans[index,k]

            xk = xhat[index,k]
            xhat_[index,k1] = xk
            
            # 更新
            if i!=self.Nt-1:
                
                y_c -= H_c[index,:,k] * np.expand_dims(xk,-1)

                kmask = np.broadcast_to(np.expand_dims(k,-1),[self.batch_size,self.Nt-i]) != \
                                np.broadcast_to(range(self.Nt-i),[self.batch_size,self.Nt-i])

                H_c = np.transpose(np.reshape(np.transpose(H_c,[1,0,2])[:,kmask] 
                        ,[self.Nr,self.batch_size,-1]),[1,0,2])
                x_c = np.reshape(x_c[kmask],[self.batch_size,-1])
                ktrans = np.reshape(ktrans[kmask],[self.batch_size,-1])

        xhat_complex = xhat_
        x_complex  = x[:,0:self.Nt] + 1j*x[:,self.Nt:self.Nt*2]

        totalnum = self.batch_size * self.Nt
        rightnum = self.equal(xhat_complex, x_complex)
        errornum = totalnum - rightnum

        nodes={
            'x_complex':x_complex,
            'xhat_complex':xhat_complex,
            'totalnum':totalnum,
            'errornum':errornum,
        }

        return nodes

    def batch_matvec_mul(self, A, b):
        """
        矩阵A与矩阵b相乘，其中A.shape=(batch_size, Nr, Nt)
        b.shape = (batch_size, Nt)
        输出矩阵C，C.shape = (batch_size, Nr)
        """
        C = np.matmul(A, np.expand_dims(b, axis=2))

        return np.squeeze(C, -1)

    def slicer_c(self, x):
        """
        软符号转换为硬符号
        Input:
        x: 检测器检测后的信号 Tensor(shape=(batchsize, Nt), dtype=complex64)
        Output:
        xhat_complex: 符号硬估计，Tensor(shape=(batchsize, Nt)， dtype=tf.int32)
        """
        Nt = np.shape(x)[1]
        x_complex = np.reshape(x,[-1, 1])
        indices = np.argmin(np.abs(x_complex - self.C), axis=1).astype(np.int32)
        xhat_complex = self.C[:,indices].T
        xhat_complex = np.reshape(xhat_complex,[-1,Nt])

        return xhat_complex
        
    def equal(self,x,y):
        """
        Computes the numbers of elements for which x and y are equal
        """
        return np.sum(np.equal(x, y).astype(np.float32))